# encoding: utf-8
"""
@author:  liaoxingyu
@contact: sherlockliao01@gmail.com
"""

import torch
import torch.nn as nn
from .center_loss import CenterLoss
from .triplet_loss import TripletLoss

def make_loss(cfg, num_classes):
    sampler = cfg.DATALOADER.SAMPLER
    feat_dim = 768 if cfg.MODEL.TRANSFORMER_TYPE in ['vit_base_patch16_224_TransReID', 'deit_base_patch16_224_TransReID'] else 384
    center_criterion = None

    if cfg.MODEL.IF_WITH_CENTER == 'yes':
        center_criterion = CenterLoss(num_classes=num_classes, feat_dim=feat_dim, use_gpu=cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available())
    else:
        print("CenterLoss is disabled (MODEL.IF_WITH_CENTER = 'no')")

    if cfg.MODEL.IF_LABELSMOOTH == 'on':
        def cross_entropy_loss(score, target):
            return nn.CrossEntropyLoss(label_smoothing=0.1)(score, target)
    else:
        def cross_entropy_loss(score, target):
            return nn.CrossEntropyLoss()(score, target)

    if sampler == 'softmax':
        def loss_func(score, feat, target, target_cam=None):
            # ============= MODIFICATION START =============
            if isinstance(score, list):
                # If score is a list (from JPM model), sum up the losses
                return sum(cross_entropy_loss(s, target) for s in score)
            else:
                # If score is a single Tensor
                return cross_entropy_loss(score, target)
            # ============== MODIFICATION END ==============

    elif sampler == 'softmax_triplet':
        # ... (your existing code here is correct, no need to change)
        triplet = TripletLoss(margin=cfg.SOLVER.MARGIN)

        def loss_func(score, feat, target, target_cam=None):
            if isinstance(score, list):
                ce_loss = sum(cross_entropy_loss(s, target) for s in score)
                trip_loss = sum(triplet(f, target)[0] for f in feat)
                return ce_loss + trip_loss
            else:
                ce_loss = cross_entropy_loss(score, target)
                trip_loss = triplet(feat, target)[0]
                return ce_loss + trip_loss

    return loss_func, center_criterion